import numpy as np 
import layers 
import funcs
from lossfuncs import LossFunc

# 神经网络
class NN:
    def __init__(self,input_shape=(1,-1),output_shape=(1,-1)):
        self.layers = list()
        self.input_shape = input_shape
        self.output_shape = output_shape

    def forward(self,x:np.ndarray)->np.ndarray:
        a = x.reshape(self.input_shape)
        for layer in self.layers:
            a = layer(a)
        return a

    def backward(self,dc_da_last:np.ndarray)->np.ndarray:
        d = dc_da_last.reshape(self.output_shape)
        for layer in self.layers[::-1]:
            d = layer.backward(d)
        return d

    def train(self,input_vec,label,loss_func:LossFunc,lr):
        y = self.forward(input_vec)
        loss = loss_func.derivate(label,y)
        self.backward(loss * -lr)

    def set_layers(self,layers):
        self.layers = layers 

    def append(self,layer):
        self.layers.append(layer)


if __name__=="__main__":
    # 使用标准的异或问题检查网络正确性
    from funcs import *
    from layers import *
    from lossfuncs import *
    my_nn = NN()
    my_nn.set_layers([
        FullConnectedLayer(2, 5),
        FuncLayer(sigmoid),
        FullConnectedLayer(5, 5),
        FuncLayer(sigmoid),
        FullConnectedLayer(5, 1),
        FuncLayer(sigmoid),
    ])
    
    data = np.array([
        [0,0],
        [0,1],
        [1,0],
        [1,1]
    ])
    label = np.array([0,1,1,0])

    epoches = 20000 

    loss_func = sse
    lr = 0.1

    for e in range(epoches):
        for i in range(len(data)):
            v = data[i]
            l = label[i]
            my_nn.train(v,l,loss_func,lr)
    for i in range(len(data)):
        v = data[i]
        l = label[i]
        predict = my_nn.forward(v)
        print(l,predict)